import argparse
from dataclasses import dataclass, field
from typing import Optional
from transformers import TrainingArguments,HfArgumentParser

# 定义模型相关的参数：
@dataclass()
class ModelArguments:
    pretrained_model_name_or_path: Optional[str] = field(
        default="CodeLlama7b_model"
    )

# 定义与数据处理相关的参数：
@dataclass()
class DataArguments:
    eval_size: Optional[int] = field(
        default=256
    )

    max_length: Optional[int] = field(
        default=512
    )

    num_data_proc: Optional[int] = field(
        default=16
    )

    train_dataset_path: Optional[str] = field(
        default="test.parquet"
    )

    eval_dataset_path: Optional[str] = field(
        default="eval.parquet"
    )


    skip_eos_token: Optional[bool] = field(
        default=False
    )

@dataclass()
class MyTrainingArguments(TrainingArguments):
    run_name: Optional[str] = field(
        default="atom"
    )

    output_dir: Optional[str] = field(
        default="output/"
    )

    per_device_train_batch_size: Optional[int] = field(
        default=4
    )

    per_device_eval_batch_size: Optional[int] = field(
        default=4
    )

    num_train_epochs: Optional[int] = field(
        default=20
    )

    weight_decay: Optional[float] = field(
        default=0
    )

    learning_rate: Optional[float] = field(
        default=1e-7
    )

    lr_scheduler_type: Optional[str] = field(
        default="cosine"
    )

    warmup_ratio: Optional[float] = field(
        default=0.1
    )

    eval_strategy: Optional[str] = field(
        default="steps"
    )

    eval_steps: Optional[int] = field(
        default=100
    )

    load_best_model_at_end: Optional[bool] = field(
        default=True
    )

    logging_strategy: Optional[str] = field(
        default="steps"
    )

    logging_steps: Optional[int] = field(
        default=1
    )

    save_strategy: Optional[str] = field(
        default="steps"
    )

    save_steps: Optional[int] = field(
        default=100
    )

    save_total_limit: Optional[int] = field(
        default=10
    )

    save_only_model: Optional[bool] = field(
        default=True
    )

    bf16: Optional[bool] = field(
        default=True
    )
    
def load_args(json_file: str):
    parser = HfArgumentParser((ModelArguments, DataArguments, MyTrainingArguments))
    # 从 JSON 文件加载配置
    model_args, data_args, training_args = parser.parse_json_file(json_file="./parser_conf_new.json")
    return model_args, data_args, training_args


